import pygame
import sys
import socket
import time
import json
import os

ZERO_FILTER = 0.07
CONFIG_FILE = "GUI_conf.json"



def input_shapers(joyaxes_lst: list) -> list:
    axis_values = joyaxes_lst.copy()
    axis_values[-2] = (-axis_values[-2] + axis_values[-1]) / 2  # Invert the Y axis
    axis_values.pop()
    axis_values = [0 if abs(value) < ZERO_FILTER else value for value in axis_values]
    
    out_values = [0., 0., 0.]
    out_values[0] = speed_map(axis_values[1], -1, 1, 100, -100)
    out_values[1] = speed_map(axis_values[3], -1, 1, 5000, -5000)
    out_values[2] = speed_map(axis_values[4], -1, 1, 5000, -5000)
    return out_values

def speed_map(value, from_min, from_max, to_min, to_max):
    """
    将数值从一个区间映射到另一个区间。

    参数:
    value (float): 要映射的数值。
    from_min (float): 原区间的最小值。
    from_max (float): 原区间的最大值。
    to_min (float): 目标区间的最小值。
    to_max (float): 目标区间的最大值。

    返回:
    float: 映射后的数值。
    """
    # 计算比例
    scale = (to_max - to_min) / (from_max - from_min)
    # 映射数值
    return to_min + (value - from_min) * scale


class JoystickGUI:
    def __init__(self):
        pygame.init()
        self.screen = pygame.display.set_mode((0,0), pygame.FULLSCREEN)
        pygame.display.set_caption("Pygame Joystick GUI")

        self.font = pygame.font.Font(None, 36)
        self.joy_id = 0
        self.ip = "0.0.0.0"
        self.port = 8080
        self.load_config()
        self.connected = False
        self.udp_socket: socket.socket = None
        self.joystick = None
        self.msg_label_str = f""
        self.udp_label_str = f"IP {self.ip}:{self.port} okay"
        self.init_pygame()
        self.input_active = None
        self.input_boxes = {
            "joy_id": pygame.Rect(200, 20, 140, 32),
            "ip": pygame.Rect(200, 60, 140, 32),
            "port": pygame.Rect(200, 100, 140, 32)
        }
        self.input_texts = {
            "joy_id": str(self.joy_id),
            "ip": self.ip,
            "port": str(self.port)
        }

        self.connect_button = pygame.Rect(200, 140, 140, 32)
        self.disconnect_button = pygame.Rect(360, 140, 140, 32)
        self.keyboard_keys = [
            {'label': '1', 'rect': pygame.Rect(20, 900, 60, 60)},
            {'label': '2', 'rect': pygame.Rect(90, 900, 60, 60)},
            {'label': '3', 'rect': pygame.Rect(160, 900, 60, 60)},
            {'label': '4', 'rect': pygame.Rect(230, 900, 60, 60)},
            {'label': '5', 'rect': pygame.Rect(300, 900, 60, 60)},
            {'label': '6', 'rect': pygame.Rect(370, 900, 60, 60)},
            {'label': '7', 'rect': pygame.Rect(440, 900, 60, 60)},
            {'label': '8', 'rect': pygame.Rect(510, 900, 60, 60)},
            {'label': '9', 'rect': pygame.Rect(580, 900, 60, 60)},
            {'label': '0', 'rect': pygame.Rect(650, 900, 60, 60)},
            {'label': '.', 'rect': pygame.Rect(720, 900, 60, 60)},
            {'label': 'Del', 'rect': pygame.Rect(790, 900, 60, 60)},
            {'label': 'Enter', 'rect': pygame.Rect(860, 900, 120, 60)},
            {'label': 'Exit',  'rect': pygame.Rect(1000, 900, 120, 60)},
        ]
    def load_config(self):
        if not os.path.exists(CONFIG_FILE):
            _config = {
                "joy_id": "0",
                "ip": "127.0.0.1",
                "port": "8080",
            }
            with open(CONFIG_FILE, 'w') as f:
                json.dump(_config, f)
        else:
            try:
                with open(CONFIG_FILE, 'r') as f:
                    _config =  json.load(f)
            except Exception as e:
                print(f"Error loading config file: {e}")
                _config = {
                    "joy_id": "0",
                    "ip": "127.0.0.1",
                    "port": "8080",
                }
                with open(CONFIG_FILE, 'w') as f:
                    json.dump(_config, f)
        self.joy_id = int(_config["joy_id"])
        self.ip = _config["ip"]
        self.port = int(_config["port"])

    def save_config(self):
        config = {
            "joy_id": self.joy_id,
            "ip": self.ip,
            "port": self.port,
        }
        with open(CONFIG_FILE, 'w') as f:
            json.dump(config, f)

    def init_pygame(self):
        pygame.joystick.init()
        if self.joystick is not None:
            self.joystick.quit()
            self.joystick = None
        try:
            self.joystick = pygame.joystick.Joystick(self.joy_id)
            self.joystick.init()
            self.msg_label_str = f"Joystick {self.joy_id} okay"
        except Exception as e:
            print(f"Error initializing joystick: {e}")
            self.msg_label_str = f"Error initializing joystick: {e}"
        

    def draw_text(self, text, pos):
        text_surface = self.font.render(text, True, (255, 255, 255))
        self.screen.blit(text_surface, pos)

    def draw_input_box(self, text, rect, active):
        color = (0, 255, 0) if active else (255, 255, 255)
        pygame.draw.rect(self.screen, color, rect, 2)
        text_surface = self.font.render(text, True, (255, 255, 255))
        self.screen.blit(text_surface, (rect.x + 5, rect.y + 5))
    

    def draw_button(self, text, rect):
        pygame.draw.rect(self.screen, (255, 255, 255), rect, 2)
        text_surface = self.font.render(text, True, (255, 255, 255))
        self.screen.blit(text_surface, (rect.x + 5, rect.y + 5))

    def draw_keyboard(self):
        for key in self.keyboard_keys:
            pygame.draw.rect(self.screen, (255, 255, 255), key['rect'], 2)
            text_surface = self.font.render(key['label'], True, (255, 255, 255))
            self.screen.blit(text_surface, (key['rect'].x + 5, key['rect'].y + 5))

    def connect_udp(self):
        if not self.connected:
            if not self.test_udp(self.ip, self.port):
                print(f"UDP connection {self.ip}:{self.port} failed")
                
                return False
            if self.udp_socket is not None:
                try:
                    self.udp_socket.close()
                except Exception as e:
                    print(f"Error closing UDP socket: {e}")
                    self.udp_label_str = f"Error closing UDP socket: {e}"
            self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.connected = True

    def disconnect_udp(self):
        if self.connected:
            self.udp_socket.close()
            self.connected = False
            self.udp_socket = None
    
    def test_udp(self, ip, udp_port):
        """
        发送 'test' 消息到指定的 UDP 端口，并检查是否收到 'okay' 响应。

        参数:
        ip (str): 目标 IP 地址。
        udp_port (int): 目标 UDP 端口。
        """
        test_message = b"test"
        expected_response = b"okay"
        timeout = 5  # 超时时间，单位为秒
        # 创建 UDP 套接字
        udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_sock.settimeout(timeout)
        try:
            # 记录发送时间
            start_time = time.time()
            # 发送 'test' 消息
            udp_sock.sendto(test_message, (ip, udp_port))
            print(f"Sent 'test' to UDP port {udp_port}")
            # 接收响应
            response, addr = udp_sock.recvfrom(1024)
            # 记录接收时间
            end_time = time.time()
            # 计算延迟
            delay = end_time - start_time
            if response == expected_response:
                print(f"Received 'okay' from {addr}")
                print(f"Round-trip delay: {delay:.6f} seconds")
                return True
            else:
                print(f"Unexpected response: {response}")
                return False
        except Exception as e:
            print(f"Erro :{e}")
            self.udp_label_str = f"Error closing UDP socket: {e}"
            return False
        finally:
            udp_sock.close()

    def run(self):
        running = True
        self.connect_udp()
        while running:
            self.screen.fill((0, 0, 0))

            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.MOUSEBUTTONDOWN:
                    ...
                elif event.type == pygame.MOUSEBUTTONUP:
                    for key, rect in self.input_boxes.items():
                        if rect.collidepoint(event.pos):
                            self.input_active = key
                            break
                    if self.connect_button.collidepoint(event.pos):
                        self.connect_udp()
                    elif self.disconnect_button.collidepoint(event.pos):
                        self.disconnect_udp()
                    
                    for key in self.keyboard_keys:
                        if key['rect'].collidepoint(event.pos) :
                            if key['label'] == 'Del' and self.input_active:
                                self.input_texts[self.input_active] = self.input_texts[self.input_active][:-1]
                            elif key['label'] == 'Enter' and self.input_active:
                                if self.input_active == "joy_id":
                                    self.joy_id = int(self.input_texts["joy_id"])
                                    self.init_pygame()
                                    self.save_config()
                                elif self.input_active == "ip":
                                    self.ip = self.input_texts["ip"]
                                    self.save_config()
                                elif self.input_active == "port":
                                    self.port = int(self.input_texts["port"])
                                    self.save_config()
                                self.input_active = None
                            
                            elif key['label'] == 'Exit':
                                running = False
                        
                            else:
                                if self.input_active:
                                    self.input_texts[self.input_active] += key['label']
                
                elif event.type == pygame.KEYDOWN:
                    if self.input_active:
                        if event.key == pygame.K_RETURN:
                            if self.input_active == "joy_id":
                                self.joy_id = int(self.input_texts["joy_id"])
                                self.init_pygame()
                            elif self.input_active == "ip":
                                self.ip = self.input_texts["ip"]
                            elif self.input_active == "port":
                                self.port = int(self.input_texts["port"])
                            self.input_active = None
                        elif event.key == pygame.K_BACKSPACE and self.input_active:
                            self.input_texts[self.input_active] = self.input_texts[self.input_active][:-1]
                        
                        elif self.input_active:
                            self.input_texts[self.input_active] += event.unicode
                    if event.key == pygame.K_ESCAPE:
                            running = False

            if self.joystick is not None:
                axis_values = [self.joystick.get_axis(i) for i in range(self.joystick.get_numaxes())]
                axis_values = [round(value, 2) for value in axis_values]
                button_values = [self.joystick.get_button(i) for i in range(self.joystick.get_numbuttons())]
                motor_speeds = input_shapers(axis_values)
            else:
                axis_values = [0] * 6
                button_values = [0] * 12
                motor_speeds = input_shapers(axis_values)
            

            self.draw_text(f"Joystick ID:", (20, 20))
            self.draw_text(f"IP Address:", (20, 60))
            self.draw_text(f"Port:", (20, 100))
            self.draw_text(f"UDP Ctrl:", (20, 140))
            self.draw_text(f"Axis values: {axis_values}", (20, 180))
            self.draw_text(f"Button values: {button_values}", (20, 220))
            self.draw_text(f"Motor speeds: {motor_speeds}", (20, 260))
            self.draw_text(self.msg_label_str, (20, 300))

            if self.connected:
                if self.joystick is not None:
                    try:
                        self.udp_socket.sendto(f":0 {motor_speeds[0]:.3f}\r\n".encode(), (self.ip, self.port))
                        self.udp_socket.sendto(f":1 {motor_speeds[1]:.3f}\r\n".encode(), (self.ip, self.port))
                        self.udp_socket.sendto(f":2 {motor_speeds[2]:.3f}\r\n".encode(), (self.ip, self.port))
                    except Exception as e:
                        self.udp_label_str = f"UDP error: {e}"
                        self.disconnect_udp()
                # self.udp_label_str = f"IP {self.ip}:{self.port} okay"
            else:
                # self.udp_label_str = f"IP {self.ip}:{self.port} error"
                ...

            for key, rect in self.input_boxes.items():
                self.draw_input_box(self.input_texts[key], rect, self.input_active == key)
            
            self.draw_text(self.udp_label_str, (20, 340))
            self.draw_button("Connect", self.connect_button)
            self.draw_button("Disconnect", self.disconnect_button)
            self.draw_keyboard()

            pygame.display.flip()
            pygame.time.wait(100)

        pygame.quit()
        sys.exit()

if __name__ == "__main__":
    gui = JoystickGUI()
    gui.run()
